import numpy as np
import time


class Ramdom:

    def __init__(self, env):
        self.env = env
        self.action, self.bandwidth, self.frequency = None, None, None
        self.cost = 1e4
        self.prob = [0.5, 0.5]

    def run(self, collector):
        for i in range(50000):
            start_timestamp = time.time()
            # 1.，模拟
            action, bandwidth, frequency = self.run_epoch(i)
            # 2. 计算当前决策的总成本
            self.env.execute(action=action, bandwidth=bandwidth, frequency=frequency)
            cost = self.env.calc_cost()
            t_ratio, cpu_ratio = self.env.calc_utilization()
            if cost < self.cost:
                self.action = action
                self.bandwidth = bandwidth
                self.frequency = frequency
                self.cost = cost
                prob = round(sum(action) / self.env.M, 2)
                self.prob = [1-prob, prob]
                print("**************************************************************************")
                print(f"Current episode: {i}")
                print(f"Offloading cost: {cost}, Offloading num: {sum(action)}, Time Ratio:{t_ratio}, CPU Ratio:{cpu_ratio}, Time:{round(time.time() - start_timestamp, 2)}")
                print(self.env.path_format())
                if collector is not None:
                    collector.collect(i, cost, sum(self.action), t_ratio, cpu_ratio)
                    collector.next()
        return self.action, self.bandwidth, self.frequency

    def run_epoch(self, i):
        action = np.zeros(self.env.M, dtype=int)
        bandwidth = np.zeros(self.env.M, dtype=int)
        frequency = np.zeros(self.env.M, dtype=int)
        for part, ues_idx in enumerate(self.env.parts_ues):
            ues_num = len(ues_idx)
            action_part = np.random.choice([0, 1], size=ues_num, p=self.prob)
            bandwidth_part, frequency_part = self.ra(ues_idx, action_part)
            timeout_idx = self.check_timeout(ues_idx, action_part, bandwidth_part, frequency_part)
            it = 1
            while len(timeout_idx) != 0:
                action_part = np.random.choice([0, 1], ues_num)
                bandwidth_part, frequency_part = self.ra(ues_idx, action_part)
                timeout_idx = self.check_timeout(ues_idx, action_part, bandwidth_part, frequency_part)
                it += 1
            for i in range(ues_num):
                ue_idx = ues_idx[i]
                action[ue_idx] = action_part[i]
                bandwidth[ue_idx] = bandwidth_part[i]
                frequency[ue_idx] = frequency_part[i]
        return action, bandwidth, frequency

    def ra(self, ues_idx, action):
        ues_num = len(ues_idx)
        bandwidth = np.zeros(ues_num, dtype=int)
        frequency = np.zeros(ues_num, dtype=int)
        off_idx = [ues_idx[i] for i in range(ues_num) if action[i]]
        meta_trans_rate = np.sqrt(self.env.meta_trans_rate)
        meta_task_cpu = np.sqrt(self.env.meta_task_cpu)
        for idx, ue_idx in enumerate(ues_idx):
            if ue_idx in off_idx:
                bandwidth[idx] = meta_trans_rate[ue_idx] * self.env.config.uav_bandwidth / sum(meta_trans_rate[off_idx])
                frequency[idx] = meta_task_cpu[ue_idx] * self.env.config.uav_cpu_frequency / sum(meta_task_cpu[off_idx])
        return bandwidth, frequency

    def check_timeout(self, ues_idx, action, bandwidth, frequency):
        with np.errstate(divide='ignore'):
            t_off = self.env.meta_trans_rate[ues_idx] / bandwidth
            t_exe = self.env.meta_task_cpu[ues_idx] / frequency
        for m in range(len(ues_idx)):
            if action[m] == 0:
                t_off[m] = 0
                t_exe[m] = 0
        t_serv = t_off + t_exe
        timeout_result = (t_serv > self.env.config.uav_hover_time)
        # noinspection PyTypeChecker
        return [idx for idx, val in enumerate(timeout_result) if val]
